import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

# <class YoloLayer(nn.Module)>
class YoloLayer(nn.Module):
    """Some Information about YoloLayer"""
    # <method __init__>
    def __init__(self, anchors, num_classes, image_size=(416, 416)):
        super(YoloLayer, self).__init__()
        self._anchors = anchors
        self._num_classes = num_classes
        self._nI_h = image_size[0] 
        self._nI_w = image_size[1]
    # <method __init__>
    # <method forward>
    def forward(self, x, nI_h=416, nI_w=416):
        self._nI_h = nI_h
        self._nI_w = nI_w
        # 
        nB = x.size(0)
        nG_h = x.size(2)
        nG_w = x.size(3)
        # 
        if x.size(1) != (self._num_classes + 4 + 1) * len(self._anchors):
            raise RuntimeError("Num of channels in x err.")
        # end-if
        # 
        nA = len(self._anchors)
        nC = self._num_classes
        # 
        fm_encoded = x.permute(0, 2, 3, 1).view(nB, nG_h, nG_w, nA, 4+1+nC) # [nB][nA, 4+1+nC][nG_w][nW] -> [nB][nG_w][nG_w][nA][4+1+nC]
        fm_encoded[ :, :, :, :, 0 ] =  torch.sigmoid(fm_encoded[ :, :, :, :, 0 ])
        fm_encoded[ :, :, :, :, 1 ] =  torch.sigmoid(fm_encoded[ :, :, :, :, 1 ])
        fm_encoded[ :, :, :, :, 4 ] =  torch.sigmoid(fm_encoded[ :, :, :, :, 4 ])
        fm_encoded[ :, :, :, :, 5:] =  torch.sigmoid(fm_encoded[ :, :, :, :, 5:])
        # 
        # test/eval mode ...
        fm_decoded = fm_encoded.clone()
        # 
        grid_h, grid_w = torch.meshgrid([torch.arange(nG_h), torch.arange(nG_w)])
        grid_h = grid_h.view(1, nG_h, nG_w).float().to(fm_encoded.device)
        grid_w = grid_w.view(1, nG_h, nG_w).float().to(fm_encoded.device)
        # 
        for na in range(nA):
            fm_decoded[ :, :, :, na, 0 ] = (fm_decoded[ :, :, :, na, 0 ] + grid_h) / (nG_h-1) * (self._nI_h-1)
            fm_decoded[ :, :, :, na, 1 ] = (fm_decoded[ :, :, :, na, 1 ] + grid_w) / (nG_w-1) * (self._nI_w-1)
            fm_decoded[ :, :, :, na, 2 ] =  torch.exp(fm_decoded[ :, :, :, na, 2 ]) * self._anchors[na][0]
            fm_decoded[ :, :, :, na, 3 ] =  torch.exp(fm_decoded[ :, :, :, na, 3 ]) * self._anchors[na][1]
        # endfor
        return fm_encoded.contiguous().view(nB, -1, 4+1+nC), fm_decoded.contiguous().view(nB, -1, 4+1+nC)
    # <method forward>
# <class YoloLayer(nn.Module)>